"""
Copyright 2020 Huawei Technologies Co., Ltd

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from xml.dom.minidom import parse
import subprocess
import argparse


def arg_parse():
    """
    Parse arguements to the operator model
    """

    parser = argparse.ArgumentParser(
        description='generate linear_transform operator model')

    parser.add_argument("-m",
                        dest='mode',
                        default='ALL',
                        type=str,
                        help="Generate mode, include combination of PCAR,\
            Flat, IVFPQ, SQ8, IVFSQ8 and ALL")

    return parser.parse_args()


def run_pcar(parameter):
    input_dim_ = parameter.getElementsByTagName("input_dim")[0]
    input_dim = int(input_dim_.childNodes[0].data)

    output_dim_ = parameter.getElementsByTagName("output_dim")[0]
    output_dim = int(output_dim_.childNodes[0].data)

    print("generate model of ID %d" % int(parameter.getAttribute("ID")))
    returnCode = subprocess.call([
        'python', 'pcar_generate_model.py',
        '-i', '%d' % (input_dim),
        '-o', '%d' % (output_dim)
    ], shell=False)
    print("")
    if returnCode != 0:
        print("ERROR: generate model of ID %d failed" %
              int(parameter.getAttribute("ID")))
        exit(1)


def run_flat(parameter):
    dim_ = parameter.getElementsByTagName("dim")[0]
    dim = int(dim_.childNodes[0].data)

    print("generate model of ID %d" % int(parameter.getAttribute("ID")))
    returnCode = subprocess.call([
        'python', 'flat_flag_generate_model.py',
        '-d', '%d' % (dim)
    ], shell=False)
    print("")
    if returnCode != 0:
        print("ERROR: generate model of ID %d failed" %
              int(parameter.getAttribute("ID")))
        exit(1)


def run_ivfpq(parameter):
    dim_ = parameter.getElementsByTagName("dim")[0]
    dim = int(dim_.childNodes[0].data)

    nlist_ = parameter.getElementsByTagName("nlist")[0]
    coarse_centroid_num = int(nlist_.childNodes[0].data)

    pq_ = parameter.getElementsByTagName("pq")[0]
    sub_quantizers_num = int(pq_.childNodes[0].data)

    print("generate model of ID %d" % int(parameter.getAttribute("ID")))
    returnCode = subprocess.call([
        'python', 'ivfpq_generate_model.py',
        '-d', '%d' % (dim),
        '-c', '%d' % (coarse_centroid_num),
        '-s', '%d' % (sub_quantizers_num)
    ], shell=False)
    print("")
    if returnCode != 0:
        print("ERROR: generate model of ID %d failed" %
              int(parameter.getAttribute("ID")))
        exit(1)


def run_sq8(parameter):
    dim_ = parameter.getElementsByTagName("dim")[0]
    dim = int(dim_.childNodes[0].data)

    print("generate model of ID %d" % int(parameter.getAttribute("ID")))
    returnCode = subprocess.call([
        'python', 'sq8_generate_model.py',
        '-d', '%d' % (dim)
    ], shell=False)
    print("")
    if returnCode != 0:
        print("ERROR: generate model of ID %d failed" %
              int(parameter.getAttribute("ID")))
        exit(1)


def run_ivfsq8(parameter):
    dim_ = parameter.getElementsByTagName("dim")[0]
    dim = int(dim_.childNodes[0].data)

    nlist_ = parameter.getElementsByTagName("nlist")[0]
    coarse_centroid_num = int(nlist_.childNodes[0].data)

    print("generate model of ID %d" % int(parameter.getAttribute("ID")))
    returnCode = subprocess.call([
        'python', 'ivfsq8_generate_model.py',
        '-d', '%d' % (dim),
        '-c', '%d' % (coarse_centroid_num)
    ], shell=False)
    print("")
    if returnCode != 0:
        print("ERROR: generate model of ID %d failed" %
              int(parameter.getAttribute("ID")))
        exit(1)


def generate_model(parameter, para_mode):
    if para_mode == 'PCAR':
        run_pcar(parameter)
    elif para_mode == 'Flat':
        run_flat(parameter)
    elif para_mode == 'IVFPQ':
        run_ivfpq(parameter)
    elif para_mode == 'SQ8':
        run_sq8(parameter)
    elif para_mode == 'IVFSQ8':
        run_ivfsq8(parameter)


if __name__ == '__main__':
    args = arg_parse()
    mode = args.mode

    modes = mode.split(',')
    for mode in modes:
        if mode not in ['PCAR', 'Flat', 'IVFPQ', 'SQ8', 'IVFSQ8', 'ALL']:
            raise RuntimeError(
                "generate mode only support PCAR, Flat, \
                    IVFPQ, SQ8, IVFSQ8 and ALL!")

    config_name = './para_table.xml'
    domTree = parse(config_name)
    rootNode = domTree.documentElement

    parameters = rootNode.getElementsByTagName("parameter")
    for parameter in parameters:
        if parameter.hasAttribute("ID"):
            mode_ = parameter.getElementsByTagName("mode")[0]
            para_mode = str(mode_.childNodes[0].data)

            if mode == 'ALL' or para_mode in modes:
                generate_model(parameter, para_mode)
            else:
                continue
